import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from scipy import signal
from scipy.io import loadmat
from scipy.stats import kurtosis, skew
import os
from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA
import warnings

warnings.filterwarnings('ignore')

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


class BearingFeatureAnalyzer:
    def __init__(self, data_path):
        self.data_path = data_path
        self.features_df = None
        self.bearing_params = {
            'SKF6205': {'n': 9, 'd': 0.3126, 'D': 1.537},  # 驱动端
            'SKF6203': {'n': 9, 'd': 0.2656, 'D': 1.122}  # 风扇端
        }

    def calculate_fault_frequencies(self, rpm, bearing_type='SKF6205'):
        """计算轴承故障特征频率"""
        fr = rpm / 60  # 转频
        params = self.bearing_params[bearing_type]
        n, d, D = params['n'], params['d'], params['D']

        # 外圈故障特征频率 BPFO
        bpfo = fr * n / 2 * (1 - d / D)
        # 内圈故障特征频率 BPFI
        bpfi = fr * n / 2 * (1 + d / D)
        # 滚动体故障特征频率 BSF
        bsf = fr * D / d * (1 - (d / D) ** 2)

        return {'BPFO': bpfo, 'BPFI': bpfi, 'BSF': bsf, 'FR': fr}

    def extract_time_domain_features(self, signal_data):
        """提取时域特征"""
        features = {}

        # 基本统计特征
        features['mean'] = np.mean(signal_data)
        features['std'] = np.std(signal_data)
        features['var'] = np.var(signal_data)
        features['rms'] = np.sqrt(np.mean(signal_data ** 2))
        features['max'] = np.max(signal_data)
        features['min'] = np.min(signal_data)
        features['peak_to_peak'] = features['max'] - features['min']

        # 高阶统计特征
        features['skewness'] = skew(signal_data)
        features['kurtosis'] = kurtosis(signal_data)

        # 形状因子
        features['shape_factor'] = features['rms'] / np.mean(np.abs(signal_data))
        features['crest_factor'] = features['max'] / features['rms']
        features['impulse_factor'] = features['max'] / np.mean(np.abs(signal_data))

        # 能量特征
        features['energy'] = np.sum(signal_data ** 2)

        return features

    def extract_frequency_domain_features(self, signal_data, fs):
        """提取频域特征"""
        # 计算FFT
        fft_data = np.fft.fft(signal_data)
        fft_magnitude = np.abs(fft_data)
        freqs = np.fft.fftfreq(len(signal_data), 1 / fs)

        # 只取正频率部分
        positive_freqs = freqs[:len(freqs) // 2]
        positive_magnitude = fft_magnitude[:len(fft_magnitude) // 2]

        features = {}

        # 频域统计特征
        features['freq_mean'] = np.mean(positive_magnitude)
        features['freq_std'] = np.std(positive_magnitude)
        features['freq_var'] = np.var(positive_magnitude)
        features['freq_max'] = np.max(positive_magnitude)

        # 频谱重心
        features['spectral_centroid'] = np.sum(positive_freqs * positive_magnitude) / np.sum(positive_magnitude)

        # 频谱峰值频率
        peak_idx = np.argmax(positive_magnitude)
        features['peak_frequency'] = positive_freqs[peak_idx]

        # 频带能量（划分为多个频带）
        total_energy = np.sum(positive_magnitude ** 2)
        freq_bands = [(0, fs / 8), (fs / 8, fs / 4), (fs / 4, 3 * fs / 8), (3 * fs / 8, fs / 2)]

        for i, (f_low, f_high) in enumerate(freq_bands):
            band_mask = (positive_freqs >= f_low) & (positive_freqs < f_high)
            band_energy = np.sum(positive_magnitude[band_mask] ** 2)
            features[f'band_{i + 1}_energy_ratio'] = band_energy / total_energy if total_energy > 0 else 0

        return features, positive_freqs, positive_magnitude

    def extract_envelope_features(self, signal_data, fs):
        """提取包络分析特征"""
        # Hilbert变换获取包络
        analytic_signal = signal.hilbert(signal_data)
        envelope = np.abs(analytic_signal)

        # 包络FFT
        envelope_fft = np.fft.fft(envelope)
        envelope_magnitude = np.abs(envelope_fft)
        envelope_freqs = np.fft.fftfreq(len(envelope), 1 / fs)

        positive_env_freqs = envelope_freqs[:len(envelope_freqs) // 2]
        positive_env_magnitude = envelope_magnitude[:len(envelope_magnitude) // 2]

        features = {}
        features['envelope_max'] = np.max(envelope)
        features['envelope_mean'] = np.mean(envelope)
        features['envelope_std'] = np.std(envelope)

        # 包络谱峰值
        env_peak_idx = np.argmax(positive_env_magnitude)
        features['envelope_peak_freq'] = positive_env_freqs[env_peak_idx]

        return features, envelope, positive_env_freqs, positive_env_magnitude

    def load_and_process_data(self, max_files_per_category=10):
        """加载并处理数据"""
        print("开始加载和处理数据...")

        all_features = []
        all_labels = []
        all_filenames = []

        # 遍历各个数据文件夹
        for root, dirs, files in os.walk(self.data_path):
            for file in files:
                if file.endswith('.mat') and len(all_features) < max_files_per_category * 4:  # 限制数据量
                    file_path = os.path.join(root, file)

                    try:
                        # 加载mat文件
                        mat_data = loadmat(file_path)

                        # 确定标签
                        if 'OR' in file or 'OR' in root:
                            label = 'OR'  # 外圈故障
                        elif 'IR' in file or 'IR' in root:
                            label = 'IR'  # 内圈故障
                        elif 'B' in file or ('B' in root and 'BA' not in root):
                            label = 'B'  # 滚动体故障
                        elif 'N' in file or 'Normal' in root:
                            label = 'N'  # 正常
                        else:
                            continue

                        # 提取振动数据 - 优先使用DE数据
                        signal_data = None
                        rpm = 1797  # 默认转速

                        for key in mat_data.keys():
                            if 'DE_time' in key:
                                signal_data = mat_data[key].flatten()
                                break
                            elif 'FE_time' in key:
                                signal_data = mat_data[key].flatten()
                                break

                        # 获取RPM信息
                        for key in mat_data.keys():
                            if 'RPM' in key:
                                rpm = float(mat_data[key][0])
                                break

                        if signal_data is not None and len(signal_data) > 1000:
                            # 截取固定长度的信号
                            signal_length = min(120000, len(signal_data))  # 约10秒数据
                            signal_data = signal_data[:signal_length]

                            # 采样频率判断
                            fs = 12000 if '12k' in root else 48000

                            # 特征提取
                            features = {}

                            # 时域特征
                            time_features = self.extract_time_domain_features(signal_data)
                            features.update(time_features)

                            # 频域特征
                            freq_features, freqs, magnitude = self.extract_frequency_domain_features(signal_data, fs)
                            features.update(freq_features)

                            # 包络特征
                            env_features, envelope, env_freqs, env_magnitude = self.extract_envelope_features(
                                signal_data, fs)
                            features.update(env_features)

                            # 故障特征频率相关特征
                            fault_freqs = self.calculate_fault_frequencies(rpm)
                            for fault_type, freq in fault_freqs.items():
                                # 在频谱中查找故障特征频率附近的能量
                                freq_tolerance = 10  # Hz
                                freq_mask = (freqs >= freq - freq_tolerance) & (freqs <= freq + freq_tolerance)
                                if np.any(freq_mask):
                                    features[f'{fault_type}_energy'] = np.sum(magnitude[freq_mask] ** 2)
                                else:
                                    features[f'{fault_type}_energy'] = 0

                            # 添加元信息
                            features['rpm'] = rpm
                            features['filename'] = file
                            features['sampling_rate'] = fs

                            all_features.append(features)
                            all_labels.append(label)
                            all_filenames.append(file)

                            print(f"处理文件: {file}, 标签: {label}, 特征数: {len(features)}")

                    except Exception as e:
                        print(f"处理文件 {file} 时出错: {e}")
                        continue

        # 转换为DataFrame
        self.features_df = pd.DataFrame(all_features)
        self.features_df['label'] = all_labels
        self.features_df['filename'] = all_filenames

        print(f"\n数据加载完成！")
        print(f"总样本数: {len(self.features_df)}")
        print(f"标签分布:\n{pd.Series(all_labels).value_counts()}")

        return self.features_df

    def visualize_signal_analysis(self, sample_files=None):
        """可视化信号分析"""
        if sample_files is None:
            # 每个类别选择一个示例文件
            sample_files = {}
            for label in ['N', 'B', 'IR', 'OR']:
                label_files = self.features_df[self.features_df['label'] == label]['filename'].tolist()
                if label_files:
                    sample_files[label] = label_files[0]

        fig, axes = plt.subplots(4, 3, figsize=(15, 12))
        fig.suptitle('轴承故障信号分析', fontsize=16)

        for i, (label, filename) in enumerate(sample_files.items()):
            if i >= 4:
                break

            # 找到文件并加载
            file_path = None
            for root, dirs, files in os.walk(self.data_path):
                if filename in files:
                    file_path = os.path.join(root, filename)
                    break

            if file_path is None:
                continue

            try:
                mat_data = loadmat(file_path)
                signal_data = None

                for key in mat_data.keys():
                    if 'DE_time' in key:
                        signal_data = mat_data[key].flatten()
                        break
                    elif 'FE_time' in key:
                        signal_data = mat_data[key].flatten()
                        break

                if signal_data is not None:
                    # 截取信号
                    signal_length = min(12000, len(signal_data))  # 1秒数据
                    signal_data = signal_data[:signal_length]
                    fs = 12000
                    time_axis = np.arange(len(signal_data)) / fs

                    # 时域图
                    axes[i, 0].plot(time_axis, signal_data)
                    axes[i, 0].set_title(f'{label} - 时域信号')
                    axes[i, 0].set_xlabel('时间 (s)')
                    axes[i, 0].set_ylabel('幅值')

                    # 频域图
                    freqs = np.fft.fftfreq(len(signal_data), 1 / fs)
                    fft_magnitude = np.abs(np.fft.fft(signal_data))
                    positive_freqs = freqs[:len(freqs) // 2]
                    positive_magnitude = fft_magnitude[:len(fft_magnitude) // 2]

                    axes[i, 1].plot(positive_freqs, positive_magnitude)
                    axes[i, 1].set_title(f'{label} - 频域信号')
                    axes[i, 1].set_xlabel('频率 (Hz)')
                    axes[i, 1].set_ylabel('幅值')
                    axes[i, 1].set_xlim(0, 2000)  # 只显示前2kHz

                    # 包络谱
                    analytic_signal = signal.hilbert(signal_data)
                    envelope = np.abs(analytic_signal)
                    env_fft = np.abs(np.fft.fft(envelope))
                    env_freqs = np.fft.fftfreq(len(envelope), 1 / fs)
                    positive_env_freqs = env_freqs[:len(env_freqs) // 2]
                    positive_env_fft = env_fft[:len(env_fft) // 2]

                    axes[i, 2].plot(positive_env_freqs, positive_env_fft)
                    axes[i, 2].set_title(f'{label} - 包络谱')
                    axes[i, 2].set_xlabel('频率 (Hz)')
                    axes[i, 2].set_ylabel('幅值')
                    axes[i, 2].set_xlim(0, 500)  # 只显示前500Hz

            except Exception as e:
                print(f"可视化文件 {filename} 时出错: {e}")

        plt.tight_layout()
        plt.show()

    def visualize_features(self):
        """特征可视化"""
        if self.features_df is None:
            print("请先加载数据！")
            return

        # 选择主要特征进行可视化
        feature_columns = [col for col in self.features_df.columns
                           if col not in ['label', 'filename', 'rpm', 'sampling_rate']]

        # 特征分布箱线图
        fig, axes = plt.subplots(3, 4, figsize=(16, 12))
        fig.suptitle('主要特征分布对比', fontsize=16)

        important_features = ['rms', 'kurtosis', 'crest_factor', 'spectral_centroid',
                              'BPFO_energy', 'BPFI_energy', 'BSF_energy', 'envelope_peak_freq',
                              'freq_max', 'peak_frequency', 'envelope_max', 'energy']

        for i, feature in enumerate(important_features):
            if i >= 12:
                break
            row, col = i // 4, i % 4

            if feature in self.features_df.columns:
                data_for_plot = []
                labels_for_plot = []

                for label in ['N', 'B', 'IR', 'OR']:
                    feature_data = self.features_df[self.features_df['label'] == label][feature]
                    if len(feature_data) > 0:
                        data_for_plot.append(feature_data)
                        labels_for_plot.append(label)

                if data_for_plot:
                    axes[row, col].boxplot(data_for_plot, labels=labels_for_plot)
                    axes[row, col].set_title(feature)
                    axes[row, col].tick_params(axis='x', rotation=45)

        plt.tight_layout()
        plt.show()

        # PCA可视化
        feature_data = self.features_df[feature_columns].fillna(0)
        scaler = StandardScaler()
        scaled_features = scaler.fit_transform(feature_data)

        pca = PCA(n_components=2)
        pca_result = pca.fit_transform(scaled_features)

        plt.figure(figsize=(10, 8))
        colors = {'N': 'green', 'B': 'blue', 'IR': 'red', 'OR': 'orange'}

        for label in colors.keys():
            mask = self.features_df['label'] == label
            if np.any(mask):
                plt.scatter(pca_result[mask, 0], pca_result[mask, 1],
                            c=colors[label], label=label, alpha=0.7)

        plt.xlabel(f'主成分1 (解释方差: {pca.explained_variance_ratio_[0]:.2%})')
        plt.ylabel(f'主成分2 (解释方差: {pca.explained_variance_ratio_[1]:.2%})')
        plt.title('特征空间PCA可视化')
        plt.legend()
        plt.grid(True, alpha=0.3)
        plt.show()

        # 特征重要性分析
        print("\nPCA特征重要性分析:")
        feature_importance = np.abs(pca.components_).mean(axis=0)
        feature_importance_df = pd.DataFrame({
            'feature': feature_columns,
            'importance': feature_importance
        }).sort_values('importance', ascending=False)

        print(feature_importance_df.head(10))

        return feature_importance_df

    def get_feature_statistics(self):
        """获取特征统计信息"""
        if self.features_df is None:
            print("请先加载数据！")
            return None

        print("\n=== 特征统计信息 ===")
        print(f"特征总数: {len([col for col in self.features_df.columns if col not in ['label', 'filename']])}")
        print(f"样本总数: {len(self.features_df)}")
        print("\n标签分布:")
        print(self.features_df['label'].value_counts())

        # 各类别的特征均值对比
        feature_columns = [col for col in self.features_df.columns
                           if col not in ['label', 'filename', 'rpm', 'sampling_rate']]

        stats_by_label = self.features_df.groupby('label')[feature_columns].mean()
        print("\n各类别主要特征均值对比:")
        important_features = ['rms', 'kurtosis', 'crest_factor', 'BPFO_energy', 'BPFI_energy', 'BSF_energy']
        for feature in important_features:
            if feature in stats_by_label.columns:
                print(f"\n{feature}:")
                print(stats_by_label[feature])

        return stats_by_label


def main():
    # 设置数据路径
    data_path = r"C:\Users\18344\Desktop\中文ACDEF题\中文赛题ACDEF题\E题\数据集\源域数据集"

    # 创建分析器
    analyzer = BearingFeatureAnalyzer(data_path)

    # 加载和处理数据
    features_df = analyzer.load_and_process_data(max_files_per_category=15)

    # 保存特征数据
    features_df.to_csv('bearing_features.csv', index=False, encoding='utf-8-sig')
    print("\n特征数据已保存到 bearing_features.csv")

    # 可视化分析
    print("\n开始可视化分析...")
    analyzer.visualize_signal_analysis()
    analyzer.visualize_features()

    # 特征统计
    stats = analyzer.get_feature_statistics()

    return analyzer, features_df


if __name__ == "__main__":
    analyzer, features_df = main()